package com.jeehentor.ai.core;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.jeehentor.ai.config.AiSessionManager;
import com.jeehentor.ai.dto.ChatMessage;
import com.jeehentor.ai.entity.AiPlatformConfig;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.MediaType;
import org.springframework.web.reactive.function.client.WebClient;
import reactor.core.publisher.Flux;

import java.util.*;

/**
 * @Description : 豆包 AI 实现类
 * @Author : wzd
 * @Date : 2025/4/23
 */
@Slf4j
@RequiredArgsConstructor
public class DoubaoClient implements AiClient {

    private final AiPlatformConfig config;
    private final WebClient webClient = WebClient.create();
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final AiSessionManager sessionManager;

//    @Override
//    public String chat(String prompt, String sessionId) {
//        try {
//            Map<String, Object> body = buildRequestWithContext(prompt, sessionId, false);
//
//            String json = webClient.post()
//                    .uri(config.getBaseUrl()+ "/api/v3/chat/completions")
//                    .header("Authorization", "Bearer " + config.getApiKey())
//                    .contentType(MediaType.APPLICATION_JSON)
//                    .bodyValue(body)
//                    .retrieve()
//                    .bodyToMono(String.class)
//                    .block();
//
//            String content = objectMapper.readTree(json)
//                    .path("output").path("text").asText();
//
//            sessionManager.appendMessage(sessionId, "user", prompt);
//            sessionManager.appendMessage(sessionId, "assistant", content);
//
//            return content;
//        } catch (Exception e) {
//            log.error("调用豆包失败", e);
//            return "AI调用失败：" + e.getMessage();
//        }
//    }
    @Override
    public String chat(String prompt, String sessionId) {
        try {
            // 构建符合官方格式的请求体
            Map<String, Object> textPart = Map.of("type", "text", "text", prompt);
            Map<String, Object> message = Map.of(
                    "role", "user",
                    "content", List.of(textPart)
            );
            Map<String, Object> message1 = Map.of(
                    "role", "system",
                    "content", "###\n" +
                            "假如你是一位专业的文章润色与配图处理师，你将根据用户提供的文章内容，来解决对文章进行润笔以及合理配置图片的任务。根据以下规则一步步执行：\n" +
                            "1. 对输入的文章进行润笔，提升文章质量。\n" +
                            "2. 根据文章内容给出配图建议，先在百度或者必应查找符合建议的图片。\n" +
                            "3. 若在百度或必应查找不到所需图片，则自己生成图片。\n" +
                            "4. 坚决不伪造假地址。\n" +
                            "5. 将润色后的文章以及图片（若有）通过Markdown格式的基本语法整理为Markdown格式输出。\n" +
                            "\n" +
                            "参考例子：\n" +
                            "无\n" +
                            "\n" +
                            "请回答问题：\n" +
                            "文章内容：XXX\n" +
                            "输出：\n" +
                            "\n" +
                            "要求：\n" +
                            "1 以Markdown格式输出\n" +
                            "2 输出内容需包含润色后的文章以及相应图片（若有）\n" +
                            "###"
            );
            ArrayList<Map<String, Object>> maps = new ArrayList<>();
            maps.add(message);
            maps.add(message1);
            Map<String, Object> body = new HashMap<>();
            body.put("model", config.getModel());
            body.put("messages", maps);

            String json = webClient.post()
                    .uri(config.getBaseUrl() + "/api/v3/chat/completions")
                    .header("Authorization", "Bearer " + config.getApiKey())
                    .contentType(MediaType.APPLICATION_JSON)
                    .bodyValue(body)
                    .retrieve()
                    .bodyToMono(String.class)
                    .block();

            // 正确解析返回内容：choices[0].message.content
            JsonNode root = objectMapper.readTree(json);
            String content = root.path("choices").path(0).path("message").path("content").asText("");

            sessionManager.appendMessage(sessionId, "user", prompt);
            sessionManager.appendMessage(sessionId, "assistant", content);

            return content;

        } catch (Exception e) {
            log.error("调用豆包失败", e);
            return "AI调用失败：" + e.getMessage();
        }
    }


    @Override
    public Flux<String> streamChat(List<ChatMessage> messages, String sessionId) {
        List<Map<String, Object>> contextList = new ArrayList<>();

        for (ChatMessage msg : messages) {
            Map<String, Object> entry = new HashMap<>();
            entry.put("role", msg.getRole());

            List<Map<String, Object>> contentList = new ArrayList<>();
            Map<String, Object> textBlock = new HashMap<>();
            textBlock.put("type", "text");
            textBlock.put("text", msg.getContent());
            contentList.add(textBlock);

            entry.put("content", contentList);
            contextList.add(entry);
        }

        Map<String, Object> req = new HashMap<>();
        req.put("model", config.getModel());
        req.put("messages", contextList);
        req.put("stream", true);

        return webClient.post()
                .uri(config.getBaseUrl() + "/api/v3/chat/completions")
                .header("Authorization", "Bearer " + config.getApiKey())
                .contentType(MediaType.APPLICATION_JSON)
                .bodyValue(req)
                .retrieve()
                .bodyToFlux(String.class)
                .flatMap(line -> {
                    if (line.startsWith("data: ")) {
                        line = line.substring(6).trim();
                    }
                    if ("[DONE]".equals(line)) return Flux.empty();

                    try {
                        JsonNode root = objectMapper.readTree(line);
                        JsonNode delta = root.path("choices").path(0).path("delta").path("content");
                        return Flux.just(delta.asText(""));
                    } catch (Exception e) {
                        log.warn("解析流式返回失败: {}", line, e);
                        return Flux.empty();
                    }
                });
    }


    private Map<String, Object> buildRequestWithContext(String prompt, String sessionId, boolean stream) {
        List<Map<String, String>> history = new ArrayList<>(sessionManager.getHistory(sessionId));
        history.add(Map.of("role", "user", "content", prompt));

        Map<String, Object> body = new HashMap<>();
        body.put("model", config.getModel());
        body.put("input", Map.of("messages", history));
        body.put("parameters", Map.of("temperature", config.getTemperature(), "incremental", stream));
        return body;
    }
}
